{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.models as models\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.transforms import RandAugment\n",
    "from torch.optim import lr_scheduler\n",
    "import numpy as np\n",
    "import time\n",
    "import copy\n",
    "import matplotlib.pyplot as plt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def get_data_loader():    \n",
    "    transform_train = transforms.Compose([\n",
    "        transforms.RandomCrop(32, padding=4),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        RandAugment(),  # 使用 RandAugment\n",
    "        transforms.ToTensor(),\n",
    "        #transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "        transforms.Normalize(\n",
    "            (0.5070751592371323, 0.48654887331495095, 0.4409178433670343),\n",
    "            (0.2673342858792401, 0.2564384629170883, 0.27615047132568404)\n",
    "        )\n",
    "    ])\n",
    "    \n",
    "    transform_test = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "    ])\n",
    "    \n",
    "\n",
    "    trainset = torchvision.datasets.CIFAR100(root='./data', train=True,\n",
    "                                             download=False, transform=transform_train)\n",
    "    trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,\n",
    "                                              shuffle=True, num_workers=2)\n",
    "\n",
    "    testset = torchvision.datasets.CIFAR100(root='./data', train=False,\n",
    "                                            download=False, transform=transform_test)\n",
    "    testloader = torch.utils.data.DataLoader(testset, batch_size=128,\n",
    "                                             shuffle=False, num_workers=2)\n",
    "        # 实例化模型\n",
    "    return trainloader,testloader\n",
    "\n",
    "\n",
    "def plot_training_results(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs):\n",
    "    # Convert PyTorch tensors to NumPy arrays\n",
    "    train_losses = np.array(train_losses)\n",
    "    train_top1_accs = np.array(train_top1_accs)\n",
    "    train_top5_accs = np.array(train_top5_accs)\n",
    "    val_losses = np.array(val_losses)\n",
    "    val_top1_accs = np.array(val_top1_accs)\n",
    "    val_top5_accs = np.array(val_top5_accs)\n",
    "\n",
    "    plt.figure(figsize=(15, 5))\n",
    "    \n",
    "    # 绘制损失曲线\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(train_losses, label='训练损失')\n",
    "    plt.plot(val_losses, label='验证损失')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.title('训练/验证损失曲线')\n",
    "    plt.legend()\n",
    "    \n",
    "    # 绘制准确率曲线\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(train_top1_accs, label='训练 Top-1 准确率')\n",
    "    plt.plot(val_top1_accs, label='验证 Top-1 准确率')\n",
    "    plt.plot(train_top5_accs, label='训练 Top-5 准确率')\n",
    "    plt.plot(val_top5_accs, label='验证 Top-5 准确率')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Accuracy')\n",
    "    plt.title('训练/验证准确率曲线')\n",
    "    plt.legend()\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "# 训练模型\n",
    "def train_model(model,trainloader,device,testloader, criterion, optimizer, scheduler, num_epochs=25):\n",
    "    since = time.time()\n",
    "       \n",
    "    train_losses = []\n",
    "    train_top1_accs = []\n",
    "    train_top5_accs = []\n",
    "    val_losses = []\n",
    "    val_top1_accs = []\n",
    "    val_top5_accs = []\n",
    "\n",
    "\n",
    "    best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    best_acc = 0.0\n",
    " \n",
    "    for epoch in range(num_epochs):\n",
    "        print('第 {}/{} 轮训练'.format(epoch, num_epochs - 1))\n",
    "        print('-' * 10)\n",
    "\n",
    "        # 每个epoch都有训练和验证阶段\n",
    "        for phase in ['train', 'val']:\n",
    "            if phase == 'train':\n",
    "                model.train()  # 设置模型为训练模式\n",
    "            else:\n",
    "                model.eval()   # 设置模型为评估模式\n",
    "\n",
    "            running_loss = 0.0\n",
    "            running_corrects = 0\n",
    "            running_top1_corrects = 0\n",
    "            running_top5_corrects = 0\n",
    "\n",
    "            # 遍历数据\n",
    "            for inputs, labels in (trainloader if phase == 'train' else testloader):\n",
    "                inputs = inputs.to(device)\n",
    "                labels = labels.to(device)\n",
    "\n",
    "                # 零化参数梯度\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # 前向传播\n",
    "                # 只在训练阶段追踪历史\n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "                    outputs = model(inputs)\n",
    "                    _, preds = torch.max(outputs, 1)\n",
    "                    loss = criterion(outputs, labels)\n",
    "\n",
    "                    # 计算 top-1 和 top-5 准确率\n",
    "                    _, top5_preds = torch.topk(outputs, 5, dim=1)\n",
    "                    top1_correct = torch.sum(preds == labels.data)\n",
    "                    top5_correct = torch.sum(top5_preds == labels.view(-1, 1))\n",
    "\n",
    "                    # 只有在训练阶段反向传播+优化\n",
    "                    if phase == 'train':\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "\n",
    "                # 统计\n",
    "                running_loss += loss.item() * inputs.size(0)\n",
    "                running_corrects += top1_correct\n",
    "                running_top1_corrects += top1_correct\n",
    "                running_top5_corrects += top5_correct\n",
    "\n",
    "            if phase == 'train':\n",
    "                scheduler.step()\n",
    "\n",
    "            epoch_loss = running_loss / len(trainloader.dataset if phase == 'train' else testloader.dataset)\n",
    "            epoch_top1_acc = running_top1_corrects.double() / len(trainloader.dataset if phase == 'train' else testloader.dataset)\n",
    "            epoch_top5_acc = running_top5_corrects.double() / len(trainloader.dataset if phase == 'train' else testloader.dataset)\n",
    "\n",
    "            print('{} 损失: {:.4f} Top-1 准确率: {:.4f} Top-5 准确率: {:.4f}'.format(\n",
    "                phase, epoch_loss, epoch_top1_acc, epoch_top5_acc))\n",
    "            \n",
    "            # 保存损失和准确率\n",
    "            if phase == 'train':\n",
    "                train_losses.append(epoch_loss)\n",
    "                train_top1_accs.append(epoch_top1_acc)\n",
    "                train_top5_accs.append(epoch_top5_acc)\n",
    "            else:\n",
    "                val_losses.append(epoch_loss)\n",
    "                val_top1_accs.append(epoch_top1_acc)\n",
    "                val_top5_accs.append(epoch_top5_acc)\n",
    "\n",
    "            # 深度复制模型\n",
    "            if phase == 'val' and epoch_top1_acc > best_acc:\n",
    "                best_acc = epoch_top1_acc\n",
    "                best_model_wts = copy.deepcopy(model.state_dict())\n",
    "\n",
    "        print()\n",
    "\n",
    "    time_elapsed = time.time() - since\n",
    "    print('训练完成，耗时 {:.0f}m {:.0f}s'.format(\n",
    "        time_elapsed // 60, time_elapsed % 60))\n",
    "    print('最佳验证集准确率: {:4f}'.format(best_acc))\n",
    "\n",
    "    # 加载最佳模型权重\n",
    "    model.load_state_dict(best_model_wts)\n",
    "    \n",
    "    # 绘制训练结果曲线\n",
    "\n",
    "    return model,(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)\n",
    " ## 修改模型\n",
    "net = ResNet18().to(device)\n",
    "num_epochs=50\n",
    " # 定义损失函数和优化器\n",
    "trainloader,testloader = get_data_loader()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)\n",
    "scheduler = lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)\n",
    "model,(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs) = train_model(net,trainloader,device,testloader, criterion, optimizer, scheduler, num_epochs)  \n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_cpu_numpy(tensor):\n",
    "    return tensor.cpu().numpy()\n",
    "def to_cpu(list_):\n",
    "    if isinstance(list_, list):\n",
    "        return list(map(to_cpu_numpy,list_))\n",
    "    else:\n",
    "        return to_cpu_numpy(list_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_training_results(train_losses, to_cpu(train_top1_accs), to_cpu(train_top5_accs), val_losses, to_cpu(val_top1_accs), to_cpu(val_top5_accs))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
